Overview:

Dynamo KVBM connector API integration with TRTLLM, including fixes after merging the connector API vLLM integration and addressing comments from #2440

Details:

The PR is based on the ongoing changes from NVIDIA/TensorRT-LLM#7228, which added TRT-LLM connector API compatibility.

Changes in this PR:

1. Added support for an external KV cache layout that is fully contiguous, since TRT-LLM’s cache layout is fully contiguous.
2. Instead of providing bytes_per_block from the KVBM leader bindings, the KVBM leader now reads bytes_per_block from the worker for a more accurate estimation. This is necessary because there is no straightforward or reliable way to extract bytes_per_block directly from the TRT-LLM KV cache.
3. Introduced additional complexity to the leader–worker barrier, requiring two synchronization steps:
   Worker → Leader: send bytes_per_block.
   Leader → Worker: send num_host_blocks and num_disk_blocks.
4. Added the basic Rust-based leader–worker integration for TRT-LLM, with some compatibility changes:
   1). When the leader calls update_state_after_alloc, no num_external_tokens is passed. For now, a HashMap is used to track each request’s num_external_tokens when calling get_num_new_matched_tokens.
   2). build_connector_metadata now uses a new implementation of apply_scheduler_output, namely apply_scheduler_output_with_computed_position, to trigger offloading. This change is required because TRT-LLM’s scheduler output does not include the number of scheduled tokens, so offloading decisions must be based on the computed position.

Issues:
A separate Python module and Rust crate will need to be built specifically for TRT-LLM. For now, TRT-LLM is included in vllm_integration to reuse some core KVBM integration code.

New changes after last review:

1. Updated the vLLM integration to read the "bytes_per_block" from the worker, instead of getting from the scheduler.
2. KVBM worker can either initialize in a blocking or non-blocking mode.
   in vLLM, the engine core will register kv cache in workers first and then initialize the leader, the engine's readiness depends on the scheduler's readiness, so the worker initialization need to be non-blocking, and the scheduler's initialization need to be blocking.
   in TRTLLM, the engine core will initialize the leader and register kv cache in workers at the same time and in the same process, he engine's readiness depends on the worker's readiness, so the worker initialization need to be blocking, and the scheduler's initialization need to be non-blocking.

Where should the reviewer start?

The reviewer could start with kvbm_connector_leader.py and kvbm_connector_worker.py. And then look into the rust bindings inside the trtllm_leader.rs and trtllm_worker.rs implementation.

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

• closes GitHub issue: #xxx